import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedKFold
from sklearn import metrics
from sklearn.ensemble import GradientBoostingClassifier
import sys
import warnings
warnings.filterwarnings('ignore')

params = {
    'n_estimators': 5000,
    'max_depth': 100,
}
model = GradientBoostingClassifier()

if len(sys.argv) > 1:
    n = int(sys.argv[1])
else:
    n = 25
trn_file = f'trn_per_{n}.csv'
print(trn_file)

trn_data = pd.read_csv(f'./data/{trn_file}', header=0)
X = trn_data.iloc[:, :-2]
y = trn_data['type']
X1 = pd.read_csv('./data/tst_desc.csv', header=0)

features = X.columns
# features = A.loc[A[1]>0.001, :][0].tolist()
# Nfea, nfea = len(A), 50
# features = A[0][Nfea-nfea:].tolist()
# print(len(features))

fold = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
models = []
pred = np.zeros((len(X1), 3))
oof = np.zeros((len(X), 3))
for index, (train_idx, val_idx) in enumerate(fold.split(X, y)):
    Xt, yt = X.loc[train_idx, features], y.iloc[train_idx]
    Xv, yv = X.loc[val_idx, features], y.iloc[val_idx]
    model = model.fit(Xt, yt)
    models.append(model)
    val_pred = model.predict_proba(Xv)
    oof[val_idx] = val_pred
    val_y = y.iloc[val_idx]
    val_pred = np.argmax(val_pred, axis=1)
    # print(index, 'val f1', metrics.f1_score(val_y, val_pred, average='macro'))

    # test_pred = model.predict_proba(X1[features])
    # pred += test_pred / 5

oof = np.argmax(oof, axis=1)
print('oof f1', metrics.f1_score(oof, y, average='macro'))
# 0.8701544575329372


import matplotlib.pyplot as plt
header = pd.DataFrame(X.columns)
feat_imp = model.feature_importances_
sort_ind = feat_imp.argsort()
A = pd.concat([header.ix[sort_ind, :].reset_index().drop('index', axis=1), pd.DataFrame(feat_imp[sort_ind])], axis=1, ignore_index=True)
plt.figure(figsize=(10, 10))
plt.barh(A[0], A[1], height=0.5)
# plt.savefig('fts/gbdt_fea_imp.png')
# plt.show()
